/**
 * https://github.com/openssl/openssl.git
 */

#include "pkcs1oaep.h"
#include "xor.h"
#include <stdlib.h>
#include <string.h>

void mgf1(const char *seed, size_t seedLen, char *maskedBuf, size_t maskLen, HashAlgo hashAlgo, size_t hLen) {
	size_t index = 0;

	size_t count = maskLen / hLen;
	if (maskLen % hLen > 0) {
		count++;
	}

	char *temp = (char *)calloc(1, seedLen + 4);
	memcpy(temp, seed, seedLen);
	for (int32_t i = 0; i < count; i++) {
		uint32_t num = (int32_t)i;

		uint8_t b[4];
		b[3] = (uint8_t)(num >> 0u);
		b[2] = (uint8_t)(num >> 8u);
		b[1] = (uint8_t)(num >> 16u);
		b[0] = (uint8_t)(num >> 24u);

		memcpy(temp + seedLen, b, 4);

		char *hash = (char *)calloc(1, hLen);
		hashAlgo(temp, seedLen + 4, hash, (int)hLen);

		for (size_t j = 0; j < hLen && index < maskLen; j++, index++) {
			maskedBuf[index] = hash[j];
		}
		free(hash);
		if (index == maskLen) {
			break;
		}
	}
	free(temp);
}

int8_t pkcs1oaepEncode(size_t keyLen,
		       const char *data,
		       size_t dataLen,
		       char *em,
		       HashAlgo hashAlgo,
		       int hashLen,
		       MGFunc mgfunc,
		       HashAlgo mgfHashAlgo,
		       int mgfHashLen,
		       RandFunc randFunc,
		       char *label,
		       size_t labelLen) {
	int8_t ret = 0;

	char *lHash = (char *)calloc(1, hashLen);
	hashAlgo(label, labelLen, lHash, hashLen);

	if (keyLen < dataLen + 2 * hashLen + 2) {
		ret = -1;
		return ret;
	}

	size_t psLen = keyLen - dataLen - 2 * hashLen - 2;

	char *ps = (char *)calloc(1, psLen);

	char *db = (char *)calloc(1, keyLen - 1 - hashLen);
	memcpy(db, lHash, hashLen);
	memcpy(db + hashLen, ps, psLen);
	db[hashLen + psLen] = 0x01;
	memcpy(db + hashLen + psLen + 1, data, dataLen);

	char *seed = (char *)calloc(1, hashLen);
	randFunc(seed, hashLen);

	char *dbMask = (char *)calloc(1, keyLen - 1 - hashLen);
	mgfunc(seed, hashLen, dbMask, keyLen - 1 - hashLen, mgfHashAlgo, mgfHashLen);

	char *maskedDB = (char *)calloc(1, keyLen - 1 - hashLen);
	xor(db, dbMask, maskedDB, keyLen - 1 - hashLen);

	char *seedMask = (char *)calloc(1, hashLen);
	mgfunc(maskedDB, keyLen - 1 - hashLen, seedMask, hashLen, mgfHashAlgo, mgfHashLen);

	char *maskedSeed = (char *)calloc(1, hashLen);
	xor(seed, seedMask, maskedSeed, hashLen);

	em[0] = 0x00;
	memcpy(em + 1, maskedSeed, hashLen);
	memcpy(em + 1 + hashLen, maskedDB, keyLen - 1 - hashLen);

	free(lHash);
	free(ps);
	free(db);
	free(seed);
	free(dbMask);
	free(maskedDB);
	free(seedMask);
	free(maskedSeed);
	return ret;
}

int8_t pkcs1oaepDecode(size_t keyLen,
               char *em,
		       char *data,
		       size_t *dataLen,
		       HashAlgo hashAlgo,
		       int hashLen,
		       MGFunc mgfunc,
		       HashAlgo mgfHashAlgo,
		       int mgfHashLen,
		       char *label,
		       size_t labelLen) {
	int8_t ret = 0;
	
	char *maskedSeed = em + 1;
	char *maskedDB = em + 1 + hashLen;

	char *seedMask = (char *)calloc(1, hashLen);
	mgfunc(maskedDB, keyLen - 1 - hashLen, seedMask, hashLen, mgfHashAlgo, mgfHashLen);

	char *seed = (char *)calloc(1, hashLen);
	xor(maskedSeed, seedMask, seed, hashLen);

	char *dbMask = (char *)calloc(1, keyLen - 1 - hashLen);
	mgfunc(seed, hashLen, dbMask, keyLen - 1 - hashLen, mgfHashAlgo, mgfHashLen);

	char *db = (char *)calloc(1, keyLen - 1 - hashLen);
	xor(maskedDB, dbMask, db, keyLen - 1 - hashLen);

	int indexOfOne = -1;
	for (int i = hashLen; i < keyLen - 1 - hashLen; i++) {
		if (db[i] == 0x01) {
			indexOfOne = i;
			break;
		}
	}

	char *temp = db + indexOfOne + 1;
	if (indexOfOne != -1) {
		memcpy(data, temp, keyLen - 1 - hashLen - indexOfOne - 1);
		*dataLen = keyLen - 1 - hashLen - indexOfOne - 1;
	} else {
		ret = -1;
	}

	free(db);
	free(seed);
	free(dbMask);
	free(seedMask);

	return ret;
}
